import os  # 导入 os 模块，用于文件和目录操作
import shutil  # 导入 shutil 模块，用于复制文件
from sklearn.model_selection import train_test_split  # 导入 train_test_split 函数，用于划分数据集

# 数据集根目录
dataset_root = 'dataset/train'  # 定义原始数据集的根目录路径

# 新创建的目录用于存放划分后的数据集
new_dataset_root = 'dataset/split_dataset'  # 定义新数据集的根目录路径

# 划分比例
train_ratio = 0.8  # 定义训练集占总数据集的比例为 80%

# 创建新目录下的 train 和 val 目录
os.makedirs(os.path.join(new_dataset_root, 'train', 'wake'), exist_ok=True)  # 创建 train/wake 子目录（如果不存在）
os.makedirs(os.path.join(new_dataset_root, 'train', 'not_wake'), exist_ok=True)  # 创建 train/not_wake 子目录
os.makedirs(os.path.join(new_dataset_root, 'val', 'wake'), exist_ok=True)  # 创建 val/wake 子目录
os.makedirs(os.path.join(new_dataset_root, 'val', 'not_wake'), exist_ok=True)  # 创建 val/not_wake 子目录


# 划分数据集
def split_dataset(class_dir):  # 定义函数 split_dataset，参数 class_dir 表示类别目录（如 wake 或 not_wake）
    # 获取所有.wav文件的完整路径
    files = [os.path.join(dataset_root, class_dir, f) for f in os.listdir(os.path.join(dataset_root, class_dir)) if
             f.endswith('.wav')]  # 遍历 class_dir 目录，获取所有 .wav 文件的完整路径

    # 划分训练集和验证集
    train_files, val_files = train_test_split(files, train_size=train_ratio, random_state=42)  # 使用 train_test_split 按比例划分数据集

    # 复制文件到新目录下的 train 和 val 文件夹
    for file in train_files:  # 遍历训练集文件
        shutil.copy(file, os.path.join(new_dataset_root, 'train', class_dir))  # 将文件复制到对应的 train 子目录

    for file in val_files:  # 遍历验证集文件
        shutil.copy(file, os.path.join(new_dataset_root, 'val', class_dir))  # 将文件复制到对应的 val 子目录


# 对 wake 和 not_wake 两个类别分别进行划分
split_dataset('wake')  # 调用 split_dataset 函数，对 wake 类别进行划分
split_dataset('not_wake')  # 调用 split_dataset 函数，对 not_wake 类别进行划分

print("数据集划分完成！")  # 打印提示信息，表示数据集划分已完成
#这段代码的主要功能是将一个包含两类数据（wake 和 not_wake）的原始数据集，按照指定比例（80% 训练集，20% 验证集）划分为训练集和验证集，并将划分后的文件分别存放到新的目录结构中